import os
import logging
import re
import time
import traceback
from typing import Dict, List, Tuple, Any, Optional, Generator, Union

try:
    import marko
    from marko.block import Heading, Paragraph, Document
    import chardet  # 添加chardet库用于检测文件编码
except ImportError:
    print("请安装必要的库: pip install marko chardet")

from toolbox import update_ui, trimmed_format_exc, CatchException, report_exception

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 最大大纲字符数限制
MAX_OUTLINE_CHARS = 1000

def extract_text_from_node(node) -> str:
    """
    从节点中提取纯文本内容，处理特殊标记
    """
    try:
        if hasattr(node, 'children'):
            if isinstance(node.children, list):
                return ''.join(extract_text_from_node(child) for child in node.children)
            else:
                text = str(node.children)
                # 清理HTML标记，如<mark>等
                text = re.sub(r'<[^>]+>', '', text)
                return text
        
        # 处理纯文本节点
        text = str(node)
        # 清理HTML标记
        text = re.sub(r'<[^>]+>', '', text)
        return text
    except Exception as e:
        logger.debug(f"提取节点文本时出错: {str(e)}")
        return ""

def extract_outline(markdown_text: str) -> Tuple[List[Dict[str, Any]], str]:
    """
    从Markdown文本中提取大纲结构，增强处理特殊格式的能力
    
    Args:
        markdown_text: Markdown格式的文本
    
    Returns:
        包含大纲结构的列表，每个元素是一个字典，包含level, title, position等信息
        以及使用的提取方法 ("marko" 或 "regex")
    """
    try:
        # 预处理markdown文本，处理特殊格式
        processed_text = markdown_text
        
        # 1. 处理HTML标记
        processed_text = re.sub(r'<mark[^>]*>(.*?)</mark>', r'\1', processed_text)
        processed_text = re.sub(r'<[^>]+>', '', processed_text)
        
        # 2. 处理图片标记，避免它们干扰标题识别
        processed_text = re.sub(r'!\[.*?\]\(.*?\)', '', processed_text)
        
        # 3. 处理数学公式
        processed_text = re.sub(r'\$.*?\$', '', processed_text)
        
        # 4. 确保代码块不干扰标题解析
        # 将代码块标记为特殊字符串，解析后再恢复
        code_blocks = []
        
        def replace_code_block(match):
            code_blocks.append(match.group(0))
            return f"CODE_BLOCK_{len(code_blocks)-1}_PLACEHOLDER"
        
        processed_text = re.sub(r'```[^`]*```', replace_code_block, processed_text)
        
        # 5. 处理mermaid图表和其他特殊块
        processed_text = re.sub(r'``mermaid[^`]*```', '', processed_text)
        processed_text = re.sub(r'``json[^`]*```', '', processed_text)
        processed_text = re.sub(r'``python[^`]*```', '', processed_text)
        
        # 6. 直接从文本中提取标题行（备用方法）
        header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
        direct_headers = []
        for match in header_pattern.finditer(processed_text):
            level = len(match.group(1))
            title = match.group(2).strip()
            line_number = processed_text[:match.start()].count('\n') + 1
            direct_headers.append({
                'level': level,
                'title': title,
                'position': (line_number, line_number)
            })
        
        # 尝试使用marko解析
        try:
            # 解析处理后的文本
            ast = marko.parse(processed_text)
            outline = []
            
            # 记录日志，帮助调试
            logger.info(f"开始提取大纲，文本长度: {len(processed_text)} 字符")
            
            for node in ast.children:
                if isinstance(node, Heading):
                    try:
                        title_text = extract_text_from_node(node)
                        # 清理标题中的特殊字符
                        title_text = title_text.strip()
                        
                        outline.append({
                            'level': node.level,
                            'title': title_text,
                            'position': (node.line_number, node.line_number + node.lines)
                        })
                        
                        logger.debug(f"提取标题: 级别={node.level}, 标题='{title_text}'")
                    except Exception as e:
                        logger.warning(f"处理标题节点时出错: {str(e)}")
                        continue
            
            # 如果marko解析失败或未找到标题，使用备用方法
            if not outline:
                logger.warning("marko解析未找到标题，使用正则表达式备用方法")
                outline = direct_headers
                return outline, "regex"
            
            logger.info(f"大纲提取完成，共找到 {len(outline)} 个标题")
            return outline, "marko"
        except Exception as e:
            logger.warning(f"marko解析失败: {str(e)}，使用正则表达式备用方法")
            return direct_headers, "regex"
            
    except Exception as e:
        error_msg = f"提取大纲时出错: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return [], "error"

def compress_outline(outline: List[Dict[str, Any]], max_chars: int, secondary_max_chars: int = 0) -> str:
    """
    压缩大纲到指定字符数
    
    Args:
        outline: 大纲结构列表
        max_chars: 最大字符数限制
        secondary_max_chars: 二级截断字符数限制，作为基础截断的补充
    
    Returns:
        压缩后的大纲文本
    """
    if not outline:
        return "无法提取大纲"
    
    # 按层级格式化大纲
    formatted_outline = []
    for item in outline:
        indent = "  " * (item['level'] - 1)
        formatted_line = f"{indent}{'#' * item['level']} {item['title']}"
        formatted_outline.append(formatted_line)
    
    full_outline = "\n".join(formatted_outline)
    
    # 如果大纲已经足够短，直接返回
    if len(full_outline) <= max_chars:
        return full_outline
    
    # 压缩策略：保留高层级标题，删除低层级标题
    compressed_outline = []
    current_level = 1
    
    while current_level <= 6:  # Markdown支持6级标题
        filtered_items = [item for item in outline if item['level'] <= current_level]
        formatted_filtered = []
        
        for item in filtered_items:
            indent = "  " * (item['level'] - 1)
            formatted_line = f"{indent}{'#' * item['level']} {item['title']}"
            formatted_filtered.append(formatted_line)
        
        compressed = "\n".join(formatted_filtered)
        
        if len(compressed) <= max_chars:
            # 添加二级截断逻辑
            if secondary_max_chars > 0 and len(compressed) < max_chars and len(full_outline) > max_chars:
                # 计算可以添加的额外字符数
                additional_chars = min(secondary_max_chars, max_chars - len(compressed))
                # 从原始大纲中截取超出部分的前additional_chars个字符
                remaining_outline = full_outline[max_chars:]
                if len(remaining_outline) > 0:
                    additional_content = remaining_outline[:additional_chars]
                    if additional_content:
                        compressed += "\n" + additional_content
            return compressed
        
        current_level += 1
    
    # 如果按层级压缩仍然超过限制，则截断
    result = full_outline[:max_chars-3] + "..."
    
    # 添加二级截断逻辑
    if secondary_max_chars > 0:
        # 从原始大纲中截取超出部分的前secondary_max_chars个字符
        remaining_outline = full_outline[max_chars-3:]  # 考虑"..."的长度
        if len(remaining_outline) > 0:
            additional_content = remaining_outline[:secondary_max_chars]
            if additional_content:
                result += "\n" + additional_content
    
    return result

def compress_outline_advanced(outline: str, max_chars: int) -> str:
    """
    压缩大纲到指定字符数，使用更智能的策略
    
    Args:
        outline: 大纲文本
        max_chars: 最大字符数限制
    
    Returns:
        压缩后的大纲文本
    """
    lines = outline.splitlines()
    if not lines:
        return "无法提取大纲"
    
    # 尝试保留尽可能多的高层级标题
    compressed_lines = []
    current_level = 1
    
    while current_level <= 6:
        filtered_lines = [line for line in lines if line.startswith("  " * (current_level - 1))]
        if filtered_lines:
            compressed_lines.extend(filtered_lines)
            # 如果压缩后仍然超过限制，则截断
            if len("\n".join(compressed_lines)) > max_chars:
                return "\n".join(compressed_lines[:max_chars-3]) + "..."
            current_level += 1
        else:
            break
    
    # 如果所有标题都被压缩，则返回原始文本的开始部分
    return outline[:max_chars-3] + "..."

def get_document_outline(file_path: str, max_chars: int = 1000, secondary_max_chars: int = 0) -> Tuple[str, str]:
    """
    获取文档的大纲并压缩
    
    Args:
        file_path: Markdown文件路径
        max_chars: 最大字符数限制
        secondary_max_chars: 二级截断字符数限制
    
    Returns:
        压缩后的大纲文本和使用的提取方法
    """
    try:
        logger.info(f"开始处理文件: {file_path}")
        
        # 尝试不同的编码方式读取文件
        encodings = ['utf-8', 'gbk', 'latin-1', 'utf-16', 'ascii']
        content = None
        
        for encoding in encodings:
            try:
                with open(file_path, 'r', encoding=encoding) as f:
                    content = f.read()
                logger.info(f"成功使用 {encoding} 编码读取文件")
                break
            except UnicodeDecodeError:
                logger.warning(f"使用 {encoding} 编码读取文件失败，尝试下一种编码")
                continue
        
        if content is None:
            # 尝试二进制模式读取，然后检测编码
            try:
                with open(file_path, 'rb') as f:
                    raw_data = f.read()
                    # 尝试检测编码
                    import chardet
                    result = chardet.detect(raw_data)
                    detected_encoding = result['encoding']
                    logger.info(f"检测到文件编码: {detected_encoding}")
                    
                    if detected_encoding:
                        content = raw_data.decode(detected_encoding)
                        logger.info(f"成功使用检测到的编码 {detected_encoding} 读取文件")
                    else:
                        logger.error(f"无法检测文件编码: {file_path}")
                        return f"处理文件时出错: 无法检测文件编码，请手动指定", "error"
            except Exception as e:
                logger.error(f"尝试二进制读取文件失败: {str(e)}")
                return f"处理文件时出错: 无法读取文件，请检查文件编码和权限", "error"
        
        if content is None:
            logger.error(f"无法使用任何编码读取文件: {file_path}")
            return f"处理文件时出错: 无法读取文件，请检查文件编码", "error"
        
        # 检查文件大小并记录日志
        file_size = os.path.getsize(file_path) / 1024  # KB
        logger.info(f"文件大小: {file_size:.2f} KB, 内容长度: {len(content)} 字符")
        
        # 检查文件是否包含常见的Markdown标记
        has_headers = bool(re.search(r'^#{1,6}\s+.+$', content, re.MULTILINE))
        has_images = bool(re.search(r'!\[.*?\]\(.*?\)', content))
        has_links = bool(re.search(r'\[.*?\]\(.*?\)', content))
        has_math = bool(re.search(r'\$.*?\$', content))
        
        logger.info(f"文件特征分析: 标题={has_headers}, 图片={has_images}, 链接={has_links}, 数学公式={has_math}")
        
        # 提取大纲
        outline, extraction_method = extract_outline(content)
        
        if not outline:
            logger.warning(f"未能从文件中提取到任何标题: {file_path}")
            # 尝试查找可能的标题行，帮助调试
            potential_headings = re.findall(r'^#+\s+.+$', content, re.MULTILINE)
            if potential_headings:
                logger.info(f"发现潜在的标题行 ({len(potential_headings)} 个)，但未能正确解析:")
                for i, heading in enumerate(potential_headings[:5]):  # 只显示前5个
                    logger.info(f"  潜在标题 {i+1}: {heading}")
                if len(potential_headings) > 5:
                    logger.info(f"  ... 还有 {len(potential_headings)-5} 个潜在标题")
                
                # 尝试直接使用正则表达式提取标题
                logger.info("尝试使用直接正则表达式提取标题...")
                header_pattern = re.compile(r'^(#{1,6})\s+(.+)$', re.MULTILINE)
                direct_headers = []
                for match in header_pattern.finditer(content):
                    level = len(match.group(1))
                    title = match.group(2).strip()
                    line_number = content[:match.start()].count('\n') + 1
                    direct_headers.append({
                        'level': level,
                        'title': title,
                        'position': (line_number, line_number)
                    })
                
                if direct_headers:
                    logger.info(f"使用正则表达式成功提取 {len(direct_headers)} 个标题")
                    outline = direct_headers
                    extraction_method = "regex"
                else:
                    return "无法提取大纲：未找到有效的标题结构", "none"
            else:
                return "无法提取大纲：未找到有效的标题结构", "none"
        
        compressed_outline = compress_outline(outline, max_chars, secondary_max_chars)
        logger.info(f"成功提取大纲，共 {len(outline)} 个标题")
        
        return compressed_outline, extraction_method
    except Exception as e:
        error_msg = f"处理文件 {file_path} 时出错: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return f"处理文件时出错: {str(e)}", "error"

def prepare_outline_for_translation(file_path: str, chatbot=None, history=None, max_chars: int = 1000, secondary_max_chars: int = 0) -> Union[str, Generator]:
    """
    为翻译准备大纲信息
    
    Args:
        file_path: Markdown文件路径
        chatbot: 聊天机器人对象，用于更新UI
        history: 历史记录
        max_chars: 大纲最大字符数限制
        secondary_max_chars: 二级截断字符数限制
        
    Returns:
        格式化的大纲信息，可直接添加到翻译prompt中
    """
    try:
        # 检查文件是否存在
        if not os.path.exists(file_path):
            error_msg = f"文件不存在: {file_path}"
            logger.error(error_msg)
            if chatbot is not None:
                chatbot.append([f"提取大纲失败", error_msg])
                if history is not None:
                    yield from update_ui(chatbot=chatbot, history=history)
            return ""
        
        # 获取压缩后的大纲
        start_time = time.time()
        outline, extraction_method = get_document_outline(file_path, 1000, 0)  // 使用默认值
        elapsed_time = time.time() - start_time
        
        # 格式化为适合添加到prompt的形式
        formatted_outline = f"""
【文档大纲信息】
以下是文档的结构大纲，请在翻译时参考：

{outline}

【大纲信息结束】
"""
        
        # 更新UI，如果提供了chatbot
        if chatbot is not None:
            chatbot.append([
                f"提取大纲: {os.path.basename(file_path)}", 
                f"成功提取大纲，共 {len(formatted_outline.splitlines())} 行（原始 {len(outline.splitlines())} 行），耗时 {elapsed_time:.2f} 秒\n使用算法: {extraction_method}"
            ])
            if history is not None:
                yield from update_ui(chatbot=chatbot, history=history)
                
        return formatted_outline
    except Exception as e:
        error_msg = f"准备大纲时出错: {str(e)}\n{trimmed_format_exc()}"
        logger.error(error_msg)
        
        # 更新UI，如果提供了chatbot
        if chatbot is not None:
            chatbot.append([f"提取大纲: {os.path.basename(file_path)}", f"出错: {error_msg}"])
            if history is not None:
                yield from update_ui(chatbot=chatbot, history=history)
                
        return ""

def get_folder_outline_cache(folder_path: str) -> Dict[str, str]:
    """
    获取文件夹中所有Markdown文件的大纲缓存
    
    Args:
        folder_path: 文件夹路径
        
    Returns:
        文件路径到大纲的映射字典
    """
    outline_cache = {}
    
    try:
        # 检查文件夹是否存在
        if not os.path.exists(folder_path) or not os.path.isdir(folder_path):
            logger.error(f"文件夹不存在或不是目录: {folder_path}")
            return {}
        
        # 查找文件夹中的full.md文件
        full_md_path = os.path.join(folder_path, 'full.md')
        if os.path.exists(full_md_path):
            outline, _ = get_document_outline(full_md_path, 1000, 0)  # 使用默认值
            outline_cache[full_md_path] = outline
            return outline_cache
            
        # 如果没有full.md，查找所有.md文件
        for file_name in os.listdir(folder_path):
            if file_name.endswith('.md'):
                file_path = os.path.join(folder_path, file_name)
                outline, _ = get_document_outline(file_path, 1000, 0)  # 使用默认值
                outline_cache[file_path] = outline
                
        return outline_cache
    except Exception as e:
        logger.error(f"获取文件夹 {folder_path} 的大纲缓存时出错: {str(e)}\n{trimmed_format_exc()}")
        return {}

def get_book_outline(folder_path: str, max_outline_chars: int, secondary_max_chars: int = 0) -> Tuple[str, str]:
    """
    获取整个书籍的大纲（优先使用full.md）
    Args:
        folder_path: 书籍文件夹路径
        max_outline_chars: 大纲最大字符数
        secondary_max_chars: 二级截断字符数限制
    Returns:
        压缩后的书籍大纲文本和使用的提取方法
    """
    try:
        # 检查文件夹是否存在
        if not os.path.exists(folder_path) or not os.path.isdir(folder_path):
            logger.error(f"文件夹不存在或不是目录: {folder_path}")
            return "无法获取书籍大纲：文件夹不存在", "error"
        logger.info(f"开始处理书籍文件夹: {folder_path}")
        # 优先查找full.md文件
        full_md_path = os.path.join(folder_path, 'full.md')
        if os.path.exists(full_md_path):
            logger.info(f"找到书籍全文文件: {full_md_path}")
            outline, extraction_method = get_document_outline(full_md_path, max_outline_chars, secondary_max_chars)
            if outline and not outline.startswith("处理文件时出错") and not outline.startswith("无法提取大纲"):
                logger.info(f"成功从全文文件提取大纲")
                # 控制最大字符数
                if len(outline) > max_outline_chars:
                    # 应用二级截断逻辑
                    if secondary_max_chars > 0 and len(outline) > max_outline_chars:
                        # 截取超出部分的前secondary_max_chars个字符
                        remaining_outline = outline[max_outline_chars:]
                        if len(remaining_outline) > 0:
                            additional_content = remaining_outline[:secondary_max_chars]
                            if additional_content:
                                outline = outline[:max_outline_chars-3] + "...\n" + additional_content
                            else:
                                outline = outline[:max_outline_chars-3] + "..."
                        else:
                            outline = outline[:max_outline_chars-3] + "..."
                    else:
                        outline = outline[:max_outline_chars-3] + "..."
                return outline, extraction_method
            else:
                logger.warning(f"从全文文件提取大纲失败，将尝试合并所有.md文件的大纲")
        # 如果没有full.md或提取失败，合并所有.md文件的大纲
        logger.info(f"将合并所有.md文件的大纲")
        all_outlines = []
        extraction_methods = []  # 记录所有使用的提取方法
        
        # 查找所有.md文件并按名称排序
        try:
            md_files = [f for f in os.listdir(folder_path) if f.endswith('.md') and not f.endswith('-中文翻译.md')]
            logger.info(f"找到 {len(md_files)} 个Markdown文件")
            if not md_files:
                logger.warning(f"文件夹中没有找到.md文件: {folder_path}")
                return "无法获取书籍大纲：文件夹中没有Markdown文件", "none"
        except Exception as e:
            logger.error(f"列出文件夹内容时出错: {str(e)}\n{traceback.format_exc()}")
            return f"无法获取书籍大纲: 列出文件夹内容时出错 - {str(e)}", "error"
        def extract_number(filename):
            import re
            match = re.search(r'(\d+)', filename)
            if match:
                return int(match.group(1))
            return 0
        try:
            md_files.sort(key=extract_number)  # 尝试按数字排序
            logger.info("成功按数字顺序排序文件")
        except:
            md_files.sort()  # 如果失败，按文件名排序
            logger.info("按文件名排序文件")
        for file_name in md_files:
            file_path = os.path.join(folder_path, file_name)
            logger.info(f"处理文件: {file_name}")
            try:
                content = None
                for encoding in ['utf-8', 'gbk', 'latin-1']:
                    try:
                        with open(file_path, 'r', encoding=encoding) as f:
                            content = f.read()
                        logger.info(f"成功使用 {encoding} 编码读取文件: {file_name}")
                        break
                    except UnicodeDecodeError:
                        logger.warning(f"使用 {encoding} 编码读取文件失败: {file_name}")
                        continue
                if content is None:
                    logger.error(f"无法使用任何编码读取文件: {file_name}")
                    continue
                outline, extraction_method = extract_outline(content)
                extraction_methods.append(extraction_method)
                if outline:
                    for item in outline:
                        item['file'] = file_name
                    all_outlines.extend(outline)
                    logger.info(f"从 {file_name} 提取了 {len(outline)} 个标题")
                else:
                    logger.warning(f"从 {file_name} 未提取到任何标题")
            except Exception as e:
                logger.warning(f"处理文件 {file_name} 时出错: {str(e)}")
                continue
        if not all_outlines:
            logger.warning(f"未能从任何文件中提取到大纲")
            return "无法获取书籍大纲：未找到有效的标题结构", "none"
        try:
            all_outlines.sort(key=lambda x: (x.get('file', ''), x['level'], x['position'][0]))
            logger.info("成功排序所有标题")
        except Exception as e:
            logger.error(f"排序标题时出错: {str(e)}")
            pass
        formatted_outline = []
        current_file = None
        for item in all_outlines:
            try:
                if item.get('file') != current_file:
                    current_file = item.get('file')
                    formatted_outline.append(f"【{current_file}】")
                indent = "  " * (item['level'] - 1)
                formatted_line = f"{indent}{'#' * item['level']} {item['title']}"
                formatted_outline.append(formatted_line)
            except Exception as e:
                logger.warning(f"格式化标题时出错: {str(e)}")
                continue
        combined_outline = "\n".join(formatted_outline)
        logger.info(f"成功合并书籍大纲，共 {len(formatted_outline)} 行，原始长度 {len(combined_outline)} 字符")

        # 使用 compress_outline 函数来处理大纲压缩
        compressed_outline = compress_outline(all_outlines, max_outline_chars, secondary_max_chars)
        logger.info(f"压缩后书籍大纲长度: {len(compressed_outline)} 字符")

        # 确定主要使用的提取方法
        if extraction_methods:
            # 统计使用最多的提取方法
            method_counts = {}
            for method in extraction_methods:
                method_counts[method] = method_counts.get(method, 0) + 1

            main_method = max(method_counts, key=method_counts.get)
            method_info = f"{main_method} (共{len(extraction_methods)}个文件中{method_counts[main_method]}个使用此方法)"
        else:
            method_info = "unknown"

        return compressed_outline, method_info
    except Exception as e:
        error_msg = f"获取书籍大纲时出错: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return f"无法获取书籍大纲: {str(e)}", "error"

def enhance_translation_prompt_with_book_outline(prompt: str, folder_path: str, max_outline_chars: int = 1000, secondary_max_chars: int = 0) -> str:
    """
    使用整个书籍的大纲增强翻译提示词
    Args:
        prompt: 原始翻译提示词
        folder_path: 书籍文件夹路径
        max_outline_chars: 大纲最大字符数
        secondary_max_chars: 二级截断字符数限制
    Returns:
        增强后的翻译提示词
    """
    try:
        logger.info(f"开始为书籍 {folder_path} 增强翻译提示词")
        logger.info(f"原始prompt长度: {len(prompt)}，前100字符: {prompt[:100]}")
        # 获取书籍大纲
        start_time = time.time()
        book_outline, extraction_method = get_book_outline(folder_path, max_outline_chars=max_outline_chars, secondary_max_chars=secondary_max_chars)
        extraction_time = time.time() - start_time
        logger.info(f"大纲提取耗时: {extraction_time:.2f} 秒")
        if book_outline.startswith("无法获取书籍大纲") or not book_outline or book_outline.startswith("处理文件时出错"):
            logger.warning(f"无法获取书籍 {folder_path} 的大纲: {book_outline}")
            folder_name = os.path.basename(folder_path)
            fallback_outline_section = f"""
【书籍大纲信息】
书籍名称: {folder_name}
注意: 由于技术原因，无法提取完整大纲，请在翻译过程中保持术语一致性和上下文连贯性。
重要提示: 翻译文本中不要包含或重复此大纲信息。只作为翻译的宏和脉络的参考。
"""
            # 无论哪种情况，都将大纲放在最前面
            enhanced_prompt = fallback_outline_section + "\n\n" + prompt
            logger.info("已将大纲信息放在提示词最前面")
            logger.info(f"使用简单书籍信息增强提示词，增强后长度: {len(enhanced_prompt)}")
            return enhanced_prompt
        metrics_file = record_outline_metrics(folder_path, book_outline, extraction_time, extraction_method, max_outline_chars, secondary_max_chars)
        if len(book_outline) > max_outline_chars:
            logger.info(f"大纲超过字符限制，进行压缩: {len(book_outline)} -> {max_outline_chars}")
            # 应用二级截断逻辑
            if secondary_max_chars > 0 and len(book_outline) > max_outline_chars:
                # 截取超出部分的前secondary_max_chars个字符
                remaining_outline = book_outline[max_outline_chars:]
                if len(remaining_outline) > 0:
                    additional_content = remaining_outline[:secondary_max_chars]
                    if additional_content:
                        book_outline = book_outline[:max_outline_chars-3] + "...\n" + additional_content
                    else:
                        book_outline = book_outline[:max_outline_chars-3] + "..."
                else:
                    book_outline = book_outline[:max_outline_chars-3] + "..."
            else:
                book_outline = book_outline[:max_outline_chars-3] + "..."
        outline_section = f"""
【书籍大纲信息】
以下是整本书的结构大纲，请在翻译时参考，保持术语一致性和上下文连贯性：

{book_outline}

【大纲信息结束】

重要提示: 翻译文本中不要包含或重复此大纲信息。翻译时仅翻译原文内容，不要在译文中输出大纲。仅作为翻译的宏和脉络的参考。
"""
        # 无论哪种情况，都将大纲放在最前面
        enhanced_prompt = outline_section + "\n\n" + prompt
        logger.info("已将大纲信息放在提示词最前面")
        logger.info(f"成功为书籍 {os.path.basename(folder_path)} 添加大纲信息，大纲长度：{len(book_outline)} 字符")
        logger.info(f"增强后prompt长度: {len(enhanced_prompt)}，前100字符: {enhanced_prompt[:100]}")
        return enhanced_prompt
    except Exception as e:
        error_msg = f"增强提示词时出错: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return prompt

@CatchException
def process_markdown_outline(txt: str, chatbot=None, history=None) -> Generator:
    """
    处理Markdown文件，提取大纲
    
    Args:
        txt: 文件路径或目录路径
        chatbot: 聊天机器人对象
        history: 历史记录
        
    Returns:
        生成器，用于更新UI
    """
    if chatbot is not None:
        chatbot.append(["Markdown大纲提取", "开始处理..."])
        if history is not None:
            yield from update_ui(chatbot=chatbot, history=history)
    
    try:
        # 检查输入是文件还是目录
        if os.path.isfile(txt):
            if txt.endswith('.md'):
                outline, extraction_method = get_document_outline(txt, 1000, 0)  // 使用默认值
                if chatbot is not None:
                    chatbot.append([f"文件: {os.path.basename(txt)}", f"大纲提取完成，共 {len(outline.splitlines())} 行\n使用算法: {extraction_method}"])
                    if history is not None:
                        yield from update_ui(chatbot=chatbot, history=history)
            else:
                if chatbot is not None:
                    chatbot.append([f"文件: {os.path.basename(txt)}", "不是Markdown文件，跳过"])
                    if history is not None:
                        yield from update_ui(chatbot=chatbot, history=history)
        elif os.path.isdir(txt):
            # 处理目录
            processed_files = 0
            for root, _, files in os.walk(txt):
                for file in files:
                    if file.endswith('.md'):
                        file_path = os.path.join(root, file)
                        outline, extraction_method = get_document_outline(file_path, 1000, 0)  // 使用默认值
                        processed_files += 1
                        if chatbot is not None:
                            chatbot.append([f"文件: {file}", f"大纲提取完成，共 {len(outline.splitlines())} 行\n使用算法: {extraction_method}"])
                            if history is not None:
                                yield from update_ui(chatbot=chatbot, history=history)
            
            if chatbot is not None:
                chatbot.append(["处理完成", f"共处理 {processed_files} 个Markdown文件"])
                if history is not None:
                    yield from update_ui(chatbot=chatbot, history=history)
        else:
            if chatbot is not None:
                chatbot.append(["错误", f"输入路径不存在: {txt}"])
                if history is not None:
                    yield from update_ui(chatbot=chatbot, history=history)
    except Exception as e:
        error_msg = f"处理出错: {str(e)}\n{trimmed_format_exc()}"
        logger.error(error_msg)
        if chatbot is not None:
            chatbot.append(["错误", error_msg])
            if history is not None:
                yield from update_ui(chatbot=chatbot, history=history)

def enhance_translation_prompt_with_outline(prompt: str, file_path: str, max_outline_chars: int = 1000) -> str:
    """
    使用文档大纲增强翻译提示词
    
    Args:
        prompt: 原始翻译提示词
        file_path: Markdown文件路径
        max_outline_chars: 大纲最大字符数限制
        
    Returns:
        增强后的翻译提示词
    """
    try:
        # 获取大纲
        outline, _ = get_document_outline(file_path, 1000, 0)  // 使用默认值
        
        # 如果大纲为空或出错，直接返回原始提示词
        if outline.startswith("处理文件时出错") or not outline:
            logger.warning(f"无法获取文件 {file_path} 的大纲，使用原始提示词")
            return prompt
        
        # 格式化大纲
        outline_section = f"""
【文档大纲信息】
以下是文档的结构大纲，请在翻译时参考：

{outline}

【大纲信息结束】
"""
        
        # 无论哪种情况，都将大纲放在最前面
        enhanced_prompt = outline_section + "\n\n" + prompt
        logger.info("已将大纲信息放在提示词最前面")
        
        return enhanced_prompt
    except Exception as e:
        logger.error(f"增强提示词时出错: {str(e)}\n{trimmed_format_exc()}")
        # 出错时返回原始提示词
        return prompt

def record_outline_metrics(folder_path: str, book_outline: str, extraction_time: float, extraction_method: str = "未知方法", max_outline_chars: int = 1000, secondary_max_chars: int = 0) -> str:
    """
    记录大纲提取的性能和质量指标
    
    Args:
        folder_path: 书籍文件夹路径
        book_outline: 提取的大纲文本
        extraction_time: 提取大纲所用的时间（秒）
        extraction_method: 使用的提取方法
        max_outline_chars: 大纲最大字符数
        secondary_max_chars: 二级截断字符数限制
        
    Returns:
        指标报告文件路径
    """
    try:
        # 创建指标文件夹
        metrics_dir = os.path.join(folder_path, "metrics")
        os.makedirs(metrics_dir, exist_ok=True)
        
        # 生成报告文件名
        report_file = os.path.join(metrics_dir, f"outline_metrics_{time.strftime('%Y%m%d%H%M%S')}.md")
        
        # 计算指标
        outline_lines = book_outline.splitlines()
        num_lines = len(outline_lines)
        chars_per_line = sum(len(line) for line in outline_lines) / max(num_lines, 1)
        
        # 分析大纲层级
        levels = {}
        for line in outline_lines:
            if line.startswith("#"):
                level = len(line) - len(line.lstrip("#"))
                levels[level] = levels.get(level, 0) + 1
        
        # 压缩大纲用于报告展示，与最终prompt中的一致
        if len(book_outline) > max_outline_chars:
            # 应用二级截断逻辑
            if secondary_max_chars > 0 and len(book_outline) > max_outline_chars:
                # 截取超出部分的前secondary_max_chars个字符
                remaining_outline = book_outline[max_outline_chars:]
                if len(remaining_outline) > 0:
                    additional_content = remaining_outline[:secondary_max_chars]
                    if additional_content:
                        compressed_outline = book_outline[:max_outline_chars-3] + "...\n" + additional_content
                    else:
                        compressed_outline = book_outline[:max_outline_chars-3] + "..."
                else:
                    compressed_outline = book_outline[:max_outline_chars-3] + "..."
            else:
                compressed_outline = book_outline[:max_outline_chars-3] + "..."
        else:
            compressed_outline = book_outline
        
        # 生成报告内容
        report_content = f"""# 大纲提取指标报告

## 基本信息
- 书籍目录: {os.path.basename(folder_path)}
- 提取时间: {time.strftime('%Y-%m-%d %H:%M:%S')}
- 处理耗时: {extraction_time:.2f} 秒
- 使用的提取算法: {extraction_method}

## 大纲统计
- 总行数: {num_lines}
- 总字符数: {len(book_outline)}
- 平均每行字符数: {chars_per_line:.2f}

## 层级分布
"""
        for level, count in sorted(levels.items()):
            report_content += f"- 第{level}级标题: {count} 个\n"
        
        report_content += f"""
## 大纲内容预览
```
{compressed_outline}
```
"""
        
        # 写入报告文件
        with open(report_file, "w", encoding="utf-8") as f:
            f.write(report_content)
        
        logger.info(f"大纲指标报告已保存到: {report_file}")
        return report_file
    except Exception as e:
        logger.error(f"记录大纲指标时出错: {str(e)}\n{traceback.format_exc()}")
        return ""

def remove_suspected_outline_blocks_in_slices(sp_file_result):
    """
    检查每个切片，找出连续（空行也算连续）七行及以上以数字开头的句子块，删除这些块，并记录被删内容。
    数字前面可以有0-5个字母（不区分大小写）、空格或#。
    返回：处理后的切片列表、被删内容字典（切片序号->被删文本列表）
    """
    import re
    new_result = []
    deleted_blocks = {}
    # 新增正则：允许0-5个字母（不区分大小写）、空格或#，其后紧跟数字
    pattern = re.compile(r'^[ \t#a-zA-Z]{0,5}\d')
    for idx, txt in enumerate(sp_file_result):
        lines = txt.splitlines()
        block = []
        block_idx = []
        result_lines = []
        i = 0
        while i < len(lines):
            line = lines[i]
            if pattern.match(line):
                temp_block = [line]
                temp_idx = [i]
                j = i + 1
                while j < len(lines):
                    next_line = lines[j]
                    if next_line.strip() == '':
                        temp_block.append(next_line)
                        temp_idx.append(j)
                        j += 1
                        continue
                    if pattern.match(next_line):
                        temp_block.append(next_line)
                        temp_idx.append(j)
                        j += 1
                    else:
                        break
                num_count = sum(1 for l in temp_block if pattern.match(l))
                if num_count >= 7:
                    if idx not in deleted_blocks:
                        deleted_blocks[idx] = []
                    deleted_blocks[idx].append('\n'.join(temp_block))
                    i = j
                else:
                    result_lines.extend(temp_block)
                    i = j
            else:
                result_lines.append(line)
                i += 1
        new_result.append('\n'.join(result_lines))
    return new_result, deleted_blocks

"""
# 集成指南：将大纲提取功能集成到 Markdown_Translate.py 中

要将大纲提取功能集成到 Markdown_Translate.py 中，请按照以下步骤操作：

## 1. 导入模块

在 Markdown_Translate.py 文件的顶部添加以下导入语句：

```python
from .Markdown_Translate_Preoutline import enhance_translation_prompt_with_outline, get_document_outline, enhance_translation_prompt_with_book_outline, get_book_outline, record_outline_metrics
```

## 2. 修改 多文件翻译 函数

在 多文件翻译 函数中，找到处理文件夹的部分，添加整本书大纲提取功能：

```python
# 当前书籍子目录完整路径
book_folder = os.path.join(project_folder, folder_name)

# 提取整本书的大纲
chatbot.append([f"书籍大纲提取: {folder_name}", "正在提取整本书的大纲..."])
yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

start_time = time.time()
book_outline = get_book_outline(book_folder)
extraction_time = time.time() - start_time

outline_lines = book_outline.splitlines() if book_outline else []
chatbot.append([f"书籍大纲提取: {folder_name}", f"提取完成，共 {len(outline_lines)} 行，耗时 {extraction_time:.2f} 秒"])
yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

# 记录大纲提取指标
metrics_file = record_outline_metrics(book_folder, book_outline, extraction_time)
if metrics_file:
    promote_file_to_downloadzone(metrics_file, chatbot=chatbot)
    chatbot.append([f"大纲指标报告", f"已生成大纲指标报告: {os.path.basename(metrics_file)}"])
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面

# 使用书籍大纲增强基础提示词
book_enhanced_prompt = enhance_translation_prompt_with_book_outline(current_prompt, book_folder)
```

## 3. 修改翻译输入构建部分

将原来的单文件大纲提取替换为使用整本书的大纲：

```python
# 构建发送给大语言模型的输入内容列表，添加大纲信息
inputs_array = [book_enhanced_prompt + f"\n\n{frag}" for frag in pfg.sp_file_contents]
inputs_show_user_array = [f"翻译 {f}" for f in pfg.sp_file_tag]
sys_prompt_array = ["You are a professional academic paper translator." + plugin_kwargs.get("additional_prompt", "") for _ in range(n_split)]
```

## 4. 测试集成

完成以上修改后，测试翻译功能，确保大纲信息正确添加到翻译提示词中，并且翻译质量有所提升。

## 注意事项

1. 确保 marko 库已安装：`pip install marko`
2. 整本书大纲提取会优先使用 full.md 文件（如果存在）
3. 如果没有 full.md 文件，系统会合并所有 .md 文件的大纲
4. 大纲提取可能会增加一些处理时间，但通常不会太明显
5. 对于非常大的文档，大纲提取可能需要更多时间，请耐心等待
6. 如果大纲提取失败，系统会自动使用原始提示词进行翻译，不会中断翻译过程
7. 系统会生成大纲提取指标报告，帮助分析大纲质量
"""

# 测试用的主函数
if __name__ == "__main__":
    # 设置日志级别为DEBUG以便查看更多信息
    logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    def test_outline_extraction(file_path):
        """测试大纲提取功能"""
        print(f"测试文件: {file_path}")
        if os.path.exists(file_path):
            print("文件存在，开始提取大纲...")
            outline = get_document_outline(file_path)
            print("\n提取的大纲:")
            print(outline)
            print(f"\n大纲行数: {len(outline.splitlines())}")
            return True
        else:
            print(f"测试文件不存在: {file_path}")
            print("请修改测试文件路径后再运行")
            return False
    
    # 测试文件路径
    test_file = "path/to/your/test.md"
    
    # 如果提供了命令行参数，使用第一个参数作为测试文件路径
    import sys
    if len(sys.argv) > 1:
        test_file = sys.argv[1]
        test_outline_extraction(test_file)
    else:
        print("未提供测试文件路径，请通过命令行参数指定，例如:")
        print("python Markdown_Translate_Preoutline.py /path/to/your/markdown/file.md")
        
        # 尝试在当前目录下查找.md文件
        current_dir = os.path.dirname(os.path.abspath(__file__))
        md_files = [f for f in os.listdir(current_dir) if f.endswith('.md')]
        if md_files:
            print(f"\n在当前目录找到以下Markdown文件，将测试第一个:")
            for i, f in enumerate(md_files):
                print(f"  {i+1}. {f}")
            test_file = os.path.join(current_dir, md_files[0])
            test_outline_extraction(test_file)
        else:
            print("当前目录下未找到Markdown文件")
